from absl import app, flags
import datetime
import logging

import pandas
import numpy
import glob

from coin.strategy.mm.fastfeature.ridge_regress import get_quantile
from coin.proto.coin_order_gateway_pb2 import OrderEvent

FLAGS = flags.FLAGS


def get_summ(oedf, colprefix, columns):
  summ = {}
  summ[colprefix] = oedf.mean()

  if False:
    for column in ['fill_qty', '%s_spread' % FLAGS.exchange]:
      corr = numpy.corrcoef(oedf['ret40'], oedf[column])[0, 1]
      summ[colprefix]['corr_%s' % column] = corr
      xq, yq = get_quantile(oedf[column], oedf['ret40'], nchunks=10)
      print(colprefix, column, corr)
      print(pandas.DataFrame({'xq': xq, 'yq': yq}).transpose().to_string())
    for column in columns:
      corr = numpy.corrcoef(oedf['ret40'], oedf['diff_%s_in_bps' % column])[0, 1]
      summ[colprefix]['corr_%s_diff' % column] = corr
      xq, yq = get_quantile(oedf['diff_%s_in_bps' % column], oedf['ret40'], nchunks=10)
      print(colprefix, column, corr)
      print(pandas.DataFrame({'xq': xq, 'yq': yq}).transpose().to_string())
  return summ


def main(_):
  import matplotlib.pylab as pylab
  params = {
      'legend.fontsize': 'small',
      'axes.labelsize': 'small',
      'axes.titlesize': 'small',
      'xtick.labelsize': 'small',
      'ytick.labelsize': 'small',
      'font.size': 8
  }
  pylab.rcParams.update(params)

  for csvfilename in glob.glob("oe_stat/*align.csv"):
    if datetime.datetime.strptime(flags.FLAGS.trading_date,
                                  "%Y%m%d").strftime("%Y-%m-%d") not in csvfilename:
      continue
    print(csvfilename)
    oedf = pandas.read_csv(csvfilename)
    if len(oedf) == 0:
      continue
    columns = [column for column in oedf.columns if column.endswith("0p")]
    columns.append("%s_buyp" % FLAGS.exchange)
    columns.append("%s_sellp" % FLAGS.exchange)
    if "%s_ask0p" % FLAGS.exchange not in columns:
      continue

    firstind = oedf['%s_ask0p' % FLAGS.exchange].first_valid_index()
    lastind = oedf['%s_ask0p' % FLAGS.exchange].last_valid_index()
    oedf = oedf.iloc[firstind:lastind]

    oedf[columns] = oedf[columns].fillna(method='ffill')
    oedf['%s_spread'
         % FLAGS.exchange] = oedf['%s_ask0p' % FLAGS.exchange] - oedf['%s_bid0p' % FLAGS.exchange]
    oedf['%s_midp' % FLAGS.exchange] = 0.5 * (oedf['%s_ask0p' % FLAGS.exchange]
                                              + oedf['%s_bid0p' % FLAGS.exchange])
    filldf = oedf.loc[oedf['fill_qty'] > 0].reset_index(drop=True)
    rets = filldf['sign'][:-1].reset_index(drop=True) * (
        filldf['%s_midp' % FLAGS.exchange][1:].reset_index(drop=True)
        - filldf['fill_price'][:-1].reset_index(drop=True)) / filldf['fill_price'][:-1].reset_index(
            drop=True)

    ret40s = filldf['sign'][:-40].reset_index(
        drop=True) * (filldf['%s_midp' % FLAGS.exchange][40:].reset_index(drop=True)
                      - filldf['fill_price'][:-40].reset_index(drop=True)
                     ) / filldf['fill_price'][:-40].reset_index(drop=True)

    ret40s_duration = (filldf['timestamp'][40:].reset_index(drop=True)
                       - filldf['timestamp'][:-40].reset_index(drop=True))

    ret_duration = (filldf['timestamp'][1:].reset_index(drop=True)
                    - filldf['timestamp'][:-1].reset_index(drop=True))

    filldf['ret'] = filldf['%s_midp' % FLAGS.exchange] * 0
    filldf['ret_duration'] = filldf['%s_midp' % FLAGS.exchange] * 0
    filldf['ret40'] = filldf['%s_midp' % FLAGS.exchange] * 0
    filldf['ret40_duration'] = filldf['%s_midp' % FLAGS.exchange] * 0

    filldf['ret'][:-1] = rets * 1e4
    filldf['ret_duration'][:-1] = ret_duration * 1e-9
    filldf['ret40'][:-40] = ret40s * 1e4
    filldf['ret40_duration'][:-40] = ret40s_duration * 1e-9
    for column in columns:
      filldf['diff_%s_in_bps'
             % column] = (filldf['fill_price'] - filldf[column]) / filldf['fill_price'] * 1e4

    for column in columns:
      oedf['diff_%s_in_bps' % column] = (oedf['price'] - oedf[column]) / oedf['price'] * 1e4

    candf = oedf.loc[oedf['event_type'] == OrderEvent.CANCEL_CONFIRMED].reset_index(drop=True)
    subdf = oedf.loc[oedf['event_type'] == OrderEvent.ORDER_SUBMITTED].reset_index(drop=True)

    # print(csvfilename)
    # print(filldf[filldf['sign'] > 0].to_string())
    symbol = filldf['symbol'][0]

    summ = {}
    summ.update(get_summ(filldf[filldf['sign'] > 0], '%s_fill_buy' % symbol, columns))
    summ.update(get_summ(filldf[filldf['sign'] < 0], '%s_fill_sell' % symbol, columns))
    # print(len(filldf))
    # print(pandas.DataFrame(summ).to_string())
    summ.update(get_summ(candf[candf['sign'] > 0], '%s_can_buy' % symbol, candf.columns))
    summ.update(get_summ(candf[candf['sign'] < 0], '%s_can_sell' % symbol, candf.columns))

    summ.update(get_summ(subdf[subdf['sign'] > 0], '%s_sub_buy' % symbol, subdf.columns))
    summ.update(get_summ(subdf[subdf['sign'] < 0], '%s_sub_sell' % symbol, subdf.columns))

    pandas.options.display.float_format = '{:.2f}'.format
    print(pandas.DataFrame(summ).to_string())

    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.subplot(411)
    plt.plot(filldf[columns], linewidth=1)
    plt.plot(filldf.loc[filldf['sign'] > 0, 'fill_price'],
             'g.',
             filldf.loc[filldf['sign'] < 0, 'fill_price'],
             'r.',
             markersize=1)
    plt.fill_between(filldf.index,
                     filldf['%s_bid0p' % FLAGS.exchange],
                     filldf['%s_ask0p' % FLAGS.exchange],
                     color='r',
                     alpha=0.2)

    if False:
      plt.twinx()
      plt.plot(filldf['timestamp'].diff(20), linewidth=0.5)
      plt.title(numpy.corrcoef(filldf['timestamp'].diff(1)[1:], filldf['ret40'][1:])[0, 1])

    plt.subplot(412)

    fqwindow = 400
    subqty = oedf['qty'] * (oedf['event_type'] == OrderEvent.ORDER_SUBMITTED)
    fillqty = oedf['fill_qty']

    frat = pandas.rolling_sum(fillqty, fqwindow) / pandas.rolling_sum(subqty, fqwindow)
    frat = frat[~numpy.isnan(frat)]
    plt.plot(frat)

    plt.subplot(413)
    plt.plot((filldf['sign'] * filldf['fill_qty']).cumsum(), linewidth=0.5, color='g')

    plt.twinx()
    plt.plot((filldf['fill_qty'] * filldf['ret40']).cumsum(), linewidth=0.5)

    plt.subplot(427)
    columns2 = ["diff_%s_in_bps" % column for column in columns]
    bid_columns2 = [column for column in columns2 if column.find("bid") >= 0]
    ask_columns2 = [column for column in columns2 if column.find("ask") >= 0]
    plt.plot(filldf[bid_columns2].loc[filldf['sign'] > 0], linewidth=1)

    bid_labels = []
    for column in bid_columns2:
      bid_labels.append("%s_avg_%.2f" % (column, filldf[column].loc[filldf['sign'] > 0].mean()))
    ask_labels = []
    for column in ask_columns2:
      ask_labels.append("%s_avg_%.2f" % (column, filldf[column].loc[filldf['sign'] < 0].mean()))

    plt.legend(bid_labels, loc='best')
    plt.title("bids diff on buy")
    plt.subplot(428)
    plt.plot(filldf[ask_columns2].loc[filldf['sign'] < 0], linewidth=1)
    plt.legend(ask_labels, loc='best')
    plt.title("asks diff on buy")
    plt.savefig("%s.png" % csvfilename, dpi=400)
    plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('market_type', 'Spot', 'Spot, Futures, Margin')
  flags.DEFINE_string('exchange', '%s', 'Exchange name.')
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
